import platgo as pg


class GA(pg.Algorithm):
    """
    应用于单目标的遗传算法
    流程：1，初始化
         2，模拟二进制交叉
         3，多项式变异
         4，环境选择
    """

    type: dict = {'single': True, 'multi': False, 'many': False, 'real': True, 'binary': True, 'permutation': True,
                  "large": True, 'expensive': False, 'constrained': False, 'preference': False, 'multimodal': False,
                  'sparse': False, 'gradient': False}

    def __init__(self, maxgen: int, problem: pg.Problem) -> None:
        # TODO 未考虑约束
        super().__init__(problem=problem, maxgen=maxgen)
        self.name = 'GA'
        self.xov = pg.operators.XovSbx()  # 模拟二进制交叉
        self.mut = pg.operators.MutPol(problem)  # 多项式变异

    def go(self, N: int = None, population: pg.Population = None) -> pg.Population:
        """
         main function for Genetic Algorithm
         if population is None, generate a new population with population size
        :param N: population size
        :param population: population to be optimized
        :return:
        """
        assert N or population, "N and population can't be both None"
        if population is None:
            pop = self.problem.init_pop(N=N)
        else:
            pop = population
            self.problem.N = pop.decs.shape[0]
        self.problem.cal_obj(pop)  # 计算目标函数值

        while self.not_terminal(pop):
            p1, p2 = pg.utils.random_selection(pop)
            offspring = self.xov(pop, p1, p2)
            offspring = self.mut(offspring)
            self.problem.cal_obj(offspring)  # 计算子代种群目标函数值
            temp_pop = offspring + pop  # 合并种群
            idx = pg.utils.naive_selection(temp_pop, pop.N)
            pop = temp_pop[idx]
        return pop
